from .table_orm import get_kline_table_orm, get_session_factory, KlineDataBase
from common.variables import time_counts, TimeFrameType
import copy
from datetime import datetime, timezone
import pandas as pd
import ray
from sqlalchemy import cast, Date, distinct, func


def ms_to_datetime_no_tz(ms: int):
    try:
        return (
            datetime.fromtimestamp(ms / 1000)
            .astimezone(timezone.utc)
            .replace(tzinfo=None)
        )
    except:
        print("debug")


def x_get_factor_from_table_only_time(
    timeframe: TimeFrameType,
    start: int,
    end: int,
    symbol: str,
    factor: str,
    session=None,
):
    """
    毫秒时间戳
    """
    table: KlineDataBase = get_kline_table_orm(timeframe)
    results = pd.read_sql(
        session.query(table.xdatetime)
        .filter(
            table.xdatetime >= ms_to_datetime_no_tz(start),
            table.xdatetime <= ms_to_datetime_no_tz(end),
            table.symbol == symbol,
            table.fname == factor,
        )
        .order_by(table.xdatetime)
        .statement,
        session.connection(),
    )
    if not results.empty:
        results[table.xdatetime.name] = results[table.xdatetime.name].apply(
            lambda x: int(x.timestamp() * 1000)
        )
    return results[table.xdatetime.name].tolist()


def x_get_factor_from_table_only_time_by_date(
    timeframe: TimeFrameType,
    start: int,
    end: int,
    symbol: str,
    factor: str,
    session=None,
):
    """
    毫秒时间戳
    """
    table: KlineDataBase = get_kline_table_orm(timeframe)
    num = time_counts["1d"] // time_counts[timeframe]
    results = pd.read_sql(
        session.query(cast(table.xdatetime, Date).label("xdatetime"))
        .filter(
            table.xdatetime >= ms_to_datetime_no_tz(start),
            table.xdatetime <= ms_to_datetime_no_tz(end),
            table.symbol == symbol,
            table.fname == factor,
        )
        .groupby(cast(table.xdatetime, Date))
        .having(func.count(distinct("xdatetime")) >= num)
        .orderby("xdatetime")
        .statement,
        session.connection(),
    )
    if not results.empty:
        results[table.xdatetime.name] = results[table.xdatetime.name].apply(
            lambda x: int(x.timestamp() * 1000)
        )
    return results[table.xdatetime.name].tolist()


def x_get_factor_from_table(
    timeframe: TimeFrameType,
    start: int,
    end: int,
    symbols: list[str],
    factors: list[str],
    session=None,
):

    table: KlineDataBase = get_kline_table_orm(timeframe)
    results = pd.read_sql(
        session.query(table)
        .filter(
            table.xdatetime >= ms_to_datetime_no_tz(start),
            table.xdatetime <= ms_to_datetime_no_tz(end),
            table.symbol.in_(symbols),
            table.fname.in_(factors),
        )
        .statement,
        session.connection(),
    )
    if not results.empty:
        results[table.xdatetime.name] = results[table.xdatetime.name].apply(
            lambda x: int(x.timestamp() * 1000)
        )
    return results


def x_get_factor_from_table_by_pivot(
    timeframe: TimeFrameType,
    start: int,
    end: int,
    symbol: str,
    factors: list[str],
    session=None,
):
    orm = get_kline_table_orm(timeframe)
    df = x_get_factor_from_table(
        timeframe, start, end, [symbol], factors, session=session
    )
    result = df.pivot(
        index=orm.xdatetime.name, columns=orm.fname.name, values=orm.value.name
    )
    return result


def x_delete_factor_from_table(
    timeframe: TimeFrameType,
    start: int,
    end: int,
    symbols: list[str],
    factors: list[str],
    session=None,
):
    table: KlineDataBase = get_kline_table_orm(timeframe)
    session.query(table).filter(
        table.xdatetime >= ms_to_datetime_no_tz(start),
        table.xdatetime <= ms_to_datetime_no_tz(end),
        table.symbol.in_(symbols),
        table.fname.in_(factors),
    ).delete(synchronize_session=False)


def x_add_factor_to_table(
    factor_value: pd.DataFrame,
    timeframe: TimeFrameType,
    pre_delete: bool = False,
    start: int = None,
    end: int = None,
    symbols: list[str] = None,
    fnames: list[str] = None,
    session=None,
):
    table = get_kline_table_orm(timeframe)
    # todo: xdatetime check
    factor_value[table.xdatetime.name] = pd.to_datetime(
        factor_value[table.xdatetime.name], unit="ms"
    )
    if pre_delete:
        x_delete_factor_from_table(
            timeframe, start, end, symbols, fnames, session=session
        )
    factor_value.to_sql(
        table.__tablename__, session.connection(), if_exists="append", index=False
    )


class TableEdit:

    def __init__(self):
        self.session_factory = get_session_factory("db_test3")

    def get_factor_from_table_only_time(
        self,
        timeframe: TimeFrameType,
        start: int,
        end: int,
        symbol: str,
        factor: str,
        session=None,
    ):
        with self.session_factory() as session:
            with session.begin():
                return x_get_factor_from_table_only_time(
                    timeframe, start, end, symbol, factor, session=session
                )

    def get_factor_from_table_only_time_by_date(
        self,
        timeframe: TimeFrameType,
        start: int,
        end: int,
        symbol: str,
        factor: str,
        session=None,
    ):
        with self.session_factory() as session:
            with session.begin():
                return x_get_factor_from_table_only_time_by_date(
                    timeframe, start, end, symbol, factor, session=session
                )

    def get_factor_from_table(
        self,
        timeframe: TimeFrameType,
        start: int,
        end: int,
        symbols: list[str],
        factors: list[str],
        session=None,
    ):
        with self.session_factory() as session:
            with session.begin():
                return x_get_factor_from_table(
                    timeframe, start, end, symbols, factors, session=session
                )

    def get_factor_from_table_by_pivot(
        self,
        timeframe: TimeFrameType,
        start: int,
        end: int,
        symbol: str,
        factors: list[str],
        session=None,
    ):
        with self.session_factory() as session:
            with session.begin():
                return x_get_factor_from_table_by_pivot(
                    timeframe, start, end, symbol, factors, session=session
                )

    def delete_factor_from_table(
        self,
        timeframe: TimeFrameType,
        start: int,
        end: int,
        symbols: list[str],
        factors: list[str],
        session=None,
    ):
        with self.session_factory() as session:
            with session.begin():
                rs = x_delete_factor_from_table(
                    timeframe, start, end, symbols, factors, session=session
                )
                return rs

    def add_factor_to_table(
        self,
        factor_value: pd.DataFrame | ray.ObjectRef,
        timeframe: TimeFrameType,
        pre_delete: bool = False,
        start: int = None,
        end: int = None,
        symbols: list[str] = None,
        fnames: list[str] = None,
        session=None,
    ):
        if isinstance(factor_value, ray.ObjectRef):
            factor_value = ray.get(factor_value)
        with self.session_factory() as session:
            with session.begin():
                rs = x_add_factor_to_table(
                    factor_value,
                    timeframe,
                    pre_delete,
                    start,
                    end,
                    symbols,
                    fnames,
                    session=session,
                )
                return rs

    def get_distinct_symbols(self, timeframe: TimeFrameType):
        table: KlineDataBase = get_kline_table_orm(timeframe)
        with self.session_factory() as session:
            with session.begin():
                return pd.read_sql(
                    session.query(distinct(table.symbol)).statement,
                    session.connection(),
                )[table.symbol.name].tolist()

    def get_distinct_factors(self, timeframe: TimeFrameType):
        table: KlineDataBase = get_kline_table_orm(timeframe)
        with self.session_factory() as session:
            with session.begin():
                return pd.read_sql(
                    session.query(distinct(table.fname)).statement,
                    session.connection(),
                )[table.fname.name].tolist()

    def get_factor_corr(
        self,
        timeframe: TimeFrameType,
        start: int,
        end: int,
        symbols: list[str],
        factors: list[str],
        delay: int = 5,
    ):
        results = []
        for symbol in symbols:
            df = self.get_factor_from_table_by_pivot(
                timeframe, start, end, symbol, factors + ["close"]
            )
            consider_cols = copy.copy(df.columns)
            if len(df) <= 0:
                continue
            df_close = df["close"]
            df["inc"] = (df_close.shift(-delay) - df_close) / (df_close + 1e-9)
            df.dropna(inplace=True)
            if len(df) <= 0:
                continue
            for fname in consider_cols:
                results.append(
                    {
                        "symbol": symbol,
                        "fname": fname,
                        "corr": df[fname].corr(df["inc"]).item(),
                    }
                )
        return results


@ray.remote(max_restarts=1, max_concurrency=1)
class RemoteTableEdit(TableEdit):
    pass


def get_table_edit_obj():
    name = "table_edit"
    space = "tradecoin"
    all_namespaces = ray.util.list_named_actors(True)
    for x in all_namespaces:
        if x["name"] == name and x["namespace"] == space:
            obj = ray.get_actor(name=name, namespace=space)
            return obj
    obj = RemoteTableEdit.options(
        name=name, lifetime="detached", namespace=space
    ).remote()
    return obj


# get_table_edit_obj()
# remote_table_edit = get_table_edit_obj()
